package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	hubur "gitee.com/wudicaidou/menciis-hubur"
	logx "gitee.com/wudicaidou/menciis-logx"
	pocsuite "gitee.com/wudicaidou/menciis-pocsuite/goClient"
	"gitee.com/wudicaidou/menciis-pocx"
	"gitee.com/wudicaidou/menciis-pocx/plugins"
	"gitee.com/wudicaidou/menciis-pocx/pocbase"
	"gitee.com/wudicaidou/menciis-pocx/snet"
	reverkit "gitee.com/wudicaidou/menciis-reverkit"
	shttp "gitee.com/wudicaidou/menciis-shttp"
	"github.com/kataras/pio"
	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	cli "github.com/urfave/cli/v2"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

func RunApp() {
	app := cli.NewApp()
	app.Usage = ""
	app.Name = "pocx"
	app.Usage = "author: zema1 && mlh"
	app.Version = "0.3 beta"
	app.Description = ""
	app.Copyright = "qdcx co."
	app.HelpName = "-h"
	app.Flags = []cli.Flag{
		&cli.StringFlag{
			Name:    "target",
			Aliases: []string{"t"},
			Usage:   "Specify a target URL [e.g. -url https://example.com]",
		},
		&cli.StringFlag{
			Name:    "target-file",
			Aliases: []string{"tF"},
			Usage:   "Select a URL file for batch identification",
		},
		&cli.StringFlag{
			Name:    "poc",
			Aliases: []string{"p"},
			Usage:   "yaml poc or yaml fingerprint Or Select a yaml file path for batch identification",
		},
		&cli.StringFlag{
			Name:    "reverse-ip",
			Aliases: []string{"rip"},
			Usage:   "set up the ip will help to launch a reverse server, ex: 127.0.0.1",
		},
		&cli.StringFlag{
			Name:        "out",
			Aliases:     []string{"o"},
			Usage:       "Save the scan results to the specified file",
			DefaultText: "\"info.log\"",
		},
		&cli.IntFlag{
			Name:    "thread",
			Aliases: []string{"c"},
			Value:   30,
			Usage:   "Number of concurrent threads",
		},
		&cli.StringFlag{
			Name:    "proxy",
			Aliases: []string{"pr"},
			Usage:   "Use proxy scan, support http/socks5 protocol [e.g. --proxy socks5://127.0.0.1:1080]",
		},
		&cli.BoolFlag{
			Name:  "verbose",
			Usage: "set log level to debug",
		},
		&cli.StringFlag{
			Name:    "python-server",
			Aliases: []string{"ps"},
			Usage:   "rpc-server",
		},
	}
	app.Action = RunServer
	err := app.Run(os.Args)
	if err != nil {
		log.Fatalf("engin err: %v", err)
		return
	}
}

func main() {
	RunApp()
}

func RunServer(c *cli.Context) error {
	t := time.Now()
	logger := logx.New()
	logger.SetTimeFormat("2006/01/02 15:04:05.000")
	var pocs []pocbase.ScanPlugin
	var targets []string
	var err error
	target := c.String("target")
	targetFile := c.String("target-file")
	yaml := c.String("poc")
	out := c.String("out")
	thread := c.Int("thread")
	proxy := c.String("proxy")
	rpcServer := c.String("python-server")
	var rpcClient *pocsuite.RpcPocClient
	if rpcServer != "" {
		rpcClient, err = pocsuite.NewRpcPocClient(rpcServer, 10)
		if err != nil {
			return err
		}
	}
	if target == "" && yaml == "" && targetFile == "" {
		return fmt.Errorf("target And poc must not be empty")
	}
	if target != "" && targetFile != "" {
		return fmt.Errorf("target or targetFile cannot be enabled at the same time")
	}
	if targetFile != "" {
		// 从文件读取目标
		fi, err := os.Open(targetFile)
		defer fi.Close()
		if err != nil {
			return err
		}
		br := bufio.NewReader(fi)
		for {
			line, _, c := br.ReadLine()
			if c == io.EOF {
				break
			}
			targets = append(targets, string(line))
		}
	} else {
		targets = []string{target}
	}
	//配置yaml文件读取
	if strings.HasSuffix(yaml, ".yaml") || strings.HasSuffix(yaml, ".yml") {
		data, err := ioutil.ReadFile(yaml)
		if err != nil {
			return err
		}
		p, err := pocx.LoadSinglePOC(data, pocx.Yml)
		if err != nil {
			return err
		}
		pocs = []pocbase.ScanPlugin{p}

	} else if strings.HasSuffix(yaml, ".py") {
		data, err := ioutil.ReadFile(yaml)
		if err != nil {
			return err
		}
		p, err := pocx.LoadSinglePOC(data, pocx.Rpy)
		if err != nil {
			return err
		}
		pocs = []pocbase.ScanPlugin{p}
	} else if yaml == "cve20200769" {
		pocs = append(pocs, plugins.Newcve20200796())
	} else if yaml == "ms17010" {
		pocs = append(pocs, plugins.NewMs17010())
	} else if yaml == "ms12020" {
		pocs = append(pocs, plugins.NewMs12020())
	} else if yaml == "cve20190708" {
		pocs = append(pocs, plugins.Newcve20190708())
	} else {
		pocs, err = pocx.LoadPocsInDir(yaml)
		if err != nil {
			return err
		}
		//pocs = append(pocs, plugins.Newcve20200796(), plugins.NewMs17010(), plugins.Newcve20190708())
	}

	defer func() {
		for _, poc := range pocs {
			poc.Close()
		}
	}()

	//配置代理
	clientOptions := shttp.DefaultClientOptions()
	clientOptions.DialTimeout = 5
	clientOptions.SoloConn = true
	//clientOptions.ReadTimeout = 10
	clientOptions.EnableHTTP2 = true
	clientOptions.FailRetries = 0
	clientOptions.DisableKeepAlives = true
	clientOptions.MaxConnsPerHost = 1
	if len(proxy) > 0 {
		logger.Infof("starting proxy %v", proxy)
		clientOptions.Proxy = proxy
	}

	//配置日志输出
	if len(out) > 0 {
		pathToAccessLog := out
		//pathToAccessLog = out + pathToAccessLog
		w, err := rotatelogs.New(
			pathToAccessLog,
			rotatelogs.WithMaxAge(24*time.Hour),
			rotatelogs.WithRotationTime(time.Hour))
		if err != nil {
			return err
		}
		defer w.Close()
		logger.SetOutput(io.MultiWriter(os.Stdout, w))
	}

	if c.Bool("verbose") {
		clientOptions.Debug = true
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var reverClient *reverkit.Client
	if c.String("reverse-ip") != "" {
		logger.Info("starting a local reverse server")
		freePort, _ := hubur.GetFreePort()
		serverConfig := &reverkit.ServerConfig{
			DBFilePath: "./reverkit.db",
			Token:      hubur.RandLower(8),
			HTTPServerConfig: reverkit.HTTPServerConfig{
				Enabled:    true,
				ListenIP:   "0.0.0.0",
				ListenPort: strconv.Itoa(freePort),
				IPHeader:   "",
			},
		}
		server, err := reverkit.NewServer(serverConfig)
		if err != nil {
			return err
		}
		err = server.Start(ctx)
		if err != nil {
			return err
		}
		defer server.Close()
		reverClient, err = reverkit.NewClient(ctx, &reverkit.ClientConfig{
			Token:       serverConfig.Token,
			HTTPBaseURL: fmt.Sprintf("http://%s:%d", c.String("reverse-ip"), freePort),
		})
		if reverClient == nil && err != nil {
			return fmt.Errorf("reverse server start err:%v", err.Error())
		}
	}
	httpClient, err := shttp.NewClient(clientOptions, nil)
	if err != nil {
		return err
	}
	netClient := snet.NewClient(&snet.ClientConfig{
		DialTimeout:  5 * time.Second,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
	})

	config := &pocx.Config{
		Logger:      logger,
		HttpClient:  httpClient,
		ReverClient: reverClient,
		NetClient:   netClient,
		RpcClient:   rpcClient,
	}

	//初始化poc对象
	ctx = pocx.ContextWithConfig(ctx, config)
	for _, yamlPoc := range pocs {
		err := yamlPoc.Init(ctx)
		if err != nil {
			logger.Error(err)
			return err
		}
	}

	group := hubur.NewSizedWaitGroup(thread)
	for _, v := range targets {
		time.Sleep(time.Second * 10)
		var in pocbase.Asset
		if strings.HasPrefix(v, "http") {
			req, _ := http.NewRequest(http.MethodGet, v, nil)
			in = pocbase.NewRequestAsset(
				&pocbase.TaskMeta{
					Type: pocbase.AssetUrlType,
				},
				&shttp.Request{
					RawRequest: req,
				},
			)
		} else {
			in, err = pocbase.NewServiceAsset(v, "tcp")
			if err != nil {
				logger.Error(err)
				continue
			}
		}
		for _, yamlPoc := range pocs {
			err := group.AddWithContext(ctx)
			if err != nil {
				logger.Error(err)
				return err
			}
			yamlPoc := yamlPoc
			go func(v string) {
				t := time.Now()

				defer func() {
					group.Done()
				}()
				snapShot, res, err := yamlPoc.Scan(ctx, in.Clone())
				logger.Infof("%v-%v %v-%v", snapShot.SrcIp, snapShot.SrcPort, snapShot.DstIp, snapShot.DstPort)
				logger.Infof("checking %s using %s", in, yamlPoc.Meta().Plugin.Name)

				if err != nil {
					if strings.Contains(err.Error(), "match nil") {
						logger.Warnf(pio.Rich(fmt.Sprintf("Target => %s run finished,err:%v", v, err.Error()), pio.Cyan))
					} else {
						logger.Warnf("Target => %s Run Finished,ERR:%v Latency => %v", v, err.Error(), time.Since(t).Seconds())
					}
					return
				}
				vulEvent, ok := res.(*pocbase.VulEvent)
				if ok {
					extractedInfo, _ := json.Marshal(vulEvent.Details.ExtractedInfo)
					logger.Info(fmt.Sprintf("Target => %s; Plugin => %s; Detatil=>%#v", v, vulEvent.Plugin.Name, string(extractedInfo)))
					logger.Logf(5, pio.Rich(fmt.Sprintf("Target => %s Plugin => %s Detatil=>%#v", v, vulEvent.Plugin.Name, string(extractedInfo)), pio.Green))
					return
				}

				return
			}(v)
		}
	}

	group.Wait()
	logger.Infof("%s", pio.Rich(fmt.Sprintf("run finished  目标总数：%v个，耗时：%v", len(targets), time.Since(t)), pio.Cyan))
	return nil
}
